# 采集工具 gradio前端界面 - 实时日志版

import gradio as gr
import json
import os
import requests
from star_xhs import XHSCrawler, cookies
import time
import threading
import logging
import io
import sys
import re
import webbrowser

# 显示Python和Gradio版本，帮助调试
print(f"Python版本: {sys.version}")
print(f"Gradio版本: {gr.__version__}")

# 配置日志处理
class LogHandler:
    def __init__(self):
        # 创建一个StringIO对象来捕获日志
        self.log_stream = io.StringIO()
        # 配置日志格式和级别
        self.logger = logging.getLogger("XHS采集器")
        self.logger.setLevel(logging.INFO)
        # 清除现有的处理器
        if self.logger.handlers:
            self.logger.handlers.clear()
        # 创建一个处理器将日志写入StringIO
        self.stream_handler = logging.StreamHandler(self.log_stream)
        self.stream_handler.setLevel(logging.INFO)
        # 设置日志格式
        formatter = logging.Formatter('[%(asctime)s] %(message)s', datefmt='%H:%M:%S')
        self.stream_handler.setFormatter(formatter)
        # 添加处理器到logger
        self.logger.addHandler(self.stream_handler)
        # 新增字段记录最后修改时间
        self.last_update = time.time()
        
    def info(self, message):
        """记录信息日志"""
        # 直接写入日志，不使用print
        self.logger.info(message)
        # 同时打印到控制台，但不使用LOG前缀，避免递归
        sys.stdout.write(f"INFO: {message}\n")
        sys.stdout.flush()
        # 更新最后修改时间
        self.last_update = time.time()
        
    def error(self, message):
        """记录错误日志"""
        self.logger.error(message)
        # 同时打印到控制台，使用sys.stderr直接写入
        sys.stderr.write(f"ERROR: {message}\n")
        sys.stderr.flush()
        # 更新最后修改时间
        self.last_update = time.time()
        
    def warning(self, message):
        """记录警告日志"""
        self.logger.warning(message)
        # 同时打印到控制台，使用sys.stdout直接写入
        sys.stdout.write(f"WARNING: {message}\n")
        sys.stdout.flush()
        # 更新最后修改时间
        self.last_update = time.time()
    
    def get_logs(self):
        """获取当前所有日志"""
        self.log_stream.seek(0)
        return self.log_stream.getvalue()
    
    def clear(self):
        """清空日志"""
        self.log_stream.truncate(0)
        self.log_stream.seek(0)
        # 更新最后修改时间
        self.last_update = time.time()

class XHSCrawlerWeb:
    def __init__(self):
        self.keyword = None  # 单个关键词
        self.crawler = None
        self.is_running = False  # 控制采集线程的标志
        self.log_handler = LogHandler()  # 日志处理器
        self.collected_notes_count = 0  # 已采集笔记计数
        self.logs_content = ""  # 存储日志内容的变量
        self.original_print = print  # 保存原始print函数
        self.auto_refresh_enabled = True  # 控制自动刷新
        self.refresh_thread = None  # 日志刷新线程
        self.stop_refresh = False  # 控制刷新线程停止的标志
        self.logs_update_callback = None  # 回调函数引用
        
    def check_cookie_status(self):
        """检查Cookie状态"""
        try:
            # 直接检查star_xhs.py中的cookies变量
            if not cookies or len(cookies[0]) < 10:  # 简单检查cookie是否为空或过短
                return "Cookie无效，请检查star_xhs.py中的cookies变量"
            
            # 尝试ping小红书网站验证连接
            try:
                response = requests.get("https://www.xiaohongshu.com", timeout=5)
                if response.status_code == 200:
                    return "Cookie有效，网站连接正常"
                else:
                    return f"网站返回异常状态码: {response.status_code}"
            except Exception as e:
                return f"无法连接到小红书网站: {str(e)}"
            
        except Exception as e:
            return f"检查Cookie出错: {str(e)}"

    def add_keyword(self, keyword):
        """设置关键词"""
        keyword = keyword.strip()
        if not keyword:
            return "", "请输入关键词"
        
        self.keyword = keyword
        return keyword, f"已设置关键词: {keyword}"

    def clear_keyword(self):
        """清空关键词"""
        self.keyword = None
        return "", "已清空关键词"
    
    def get_logs(self):
        """获取当前日志，用于UI显示"""
        return self.log_handler.get_logs()

    # 重写爬虫的输出钩子，简化日志格式
    def setup_crawler_hooks(self):
        """设置爬虫的钩子函数，拦截并简化日志输出"""
        # 保存原始print函数
        self.original_print = print
        
        # 自定义print函数
        def custom_print(*args, **kwargs):
            # 将所有参数转换为一个字符串
            message = " ".join(str(arg) for arg in args)
            
            # 简化日志输出
            if "标题元素" in message or "作者元素" in message:
                return  # 跳过冗余日志
            
            elif "已采集文章" in message:
                title = message.split("：")[1] if "：" in message else "未知标题"
                self.collected_notes_count += 1
                # 使用sys.stdout直接写入而不是调用log_handler.info
                sys.stdout.write(f"INFO: 采集第{self.collected_notes_count}篇笔记: {title}\n")
                sys.stdout.flush()
                # 直接使用logger写入日志流
                self.log_handler.logger.info(f"采集第{self.collected_notes_count}篇笔记: {title}")
                return
                
            elif "评论内容" in message or "评论用户" in message or "评论时间" in message:
                return  # 跳过评论细节
                
            elif "即将保存评论" in message:
                return  # 跳过保存评论详情
                
            elif "已采集第" in message and "条评论" in message:
                try:
                    count = message.split("已采集第")[1].split("条评论")[0].strip()
                    # 使用直接写入方式
                    sys.stdout.write(f"INFO: 已采集 {count} 条评论\n")
                    sys.stdout.flush()
                    self.log_handler.logger.info(f"已采集 {count} 条评论")
                except:
                    sys.stdout.write("INFO: 已采集评论\n")
                    sys.stdout.flush()
                    self.log_handler.logger.info("已采集评论")
                return
                
            elif "笔记评论采集完成" in message:
                try:
                    count_match = re.search(r"共采集(\d+)条评论", message)
                    if count_match:
                        count = count_match.group(1)
                        sys.stdout.write(f"INFO: 评论采集完成，共 {count} 条\n")
                        sys.stdout.flush()
                        self.log_handler.logger.info(f"评论采集完成，共 {count} 条")
                    else:
                        sys.stdout.write("INFO: 评论采集完成\n")
                        sys.stdout.flush()
                        self.log_handler.logger.info("评论采集完成")
                    
                    sys.stdout.write(f"INFO: 采集第{self.collected_notes_count}篇笔记结束\n")
                    sys.stdout.flush()
                    self.log_handler.logger.info(f"采集第{self.collected_notes_count}篇笔记结束")
                except:
                    sys.stdout.write("INFO: 评论采集完成\n")
                    sys.stdout.flush()
                    self.log_handler.logger.info("评论采集完成")
                    
                    sys.stdout.write(f"INFO: 采集第{self.collected_notes_count}篇笔记结束\n")
                    sys.stdout.flush()
                    self.log_handler.logger.info(f"采集第{self.collected_notes_count}篇笔记结束")
                return
                
            # 其他关键信息保留
            elif any(keyword in message for keyword in ["登录", "搜索", "采集", "错误", "失败", "cookie", "笔记"]):
                sys.stdout.write(f"INFO: {message}\n")
                sys.stdout.flush()
                self.log_handler.logger.info(message)
            
            # 使用原始print打印到控制台（调试时可以开启）
            # self.original_print(*args, **kwargs)
        
        # 替换全局print函数
        import builtins
        builtins.print = custom_print
    
    def restore_print(self):
        """恢复原始的print函数"""
        import builtins
        builtins.print = self.original_print
        self.original_print("已恢复原始print函数")

    def crawler_thread(self, notes_count, comments_count, keyword, status_callback=None):
        """爬虫线程函数"""
        try:
            self.is_running = True
            self.collected_notes_count = 0
            self.log_handler.clear()  # 清空之前的日志
            
            # 直接写入日志流
            self.log_handler.logger.info(f"开始采集关键词: {keyword}")
            sys.stdout.write(f"INFO: 开始采集关键词: {keyword}\n")
            sys.stdout.flush()
            
            # 设置爬虫钩子，简化日志 - 在创建爬虫实例前设置
            self.setup_crawler_hooks()
            
            # 创建爬虫实例
            self.crawler = XHSCrawler()
            self.crawler.set_limits(notes_count, comments_count)
            
            # 登录
            sys.stdout.write("INFO: 正在登录小红书...\n")
            sys.stdout.flush()
            self.log_handler.logger.info("正在登录小红书...")
            
            if not self.crawler.login():
                sys.stderr.write("ERROR: 登录失败，请检查Cookie\n")
                sys.stderr.flush()
                self.log_handler.logger.error("登录失败，请检查Cookie")
                if status_callback:
                    status_callback("登录失败，请检查Cookie")
                return
            
            sys.stdout.write("INFO: 登录成功，开始执行搜索...\n")
            sys.stdout.flush()
            self.log_handler.logger.info("登录成功，开始执行搜索...")
            
            # 初始化数据采集器并设置关键词
            from collect_data import DataCollector
            self.note_collector = DataCollector()
            self.note_collector.set_keyword(keyword)
            
            # 执行搜索和采集
            if self.crawler.search_and_filter(keyword):
                sys.stdout.write(f"INFO: 已搜索关键词并设置筛选条件，开始采集笔记...\n")
                sys.stdout.flush()
                self.log_handler.logger.info(f"已搜索关键词并设置筛选条件，开始采集笔记...")
                
                if self.crawler.collect_all_notes():
                    sys.stdout.write(f"INFO: 关键词 {keyword} 采集完成\n")
                    sys.stdout.flush()
                    self.log_handler.logger.info(f"关键词 {keyword} 采集完成")
                    if status_callback:
                        status_callback(f"采集任务已完成，关键词: {keyword}")
                else:
                    sys.stderr.write(f"ERROR: 采集关键词 {keyword} 的笔记时出错\n")
                    sys.stderr.flush()
                    self.log_handler.logger.error(f"采集关键词 {keyword} 的笔记时出错")
                    if status_callback:
                        status_callback(f"采集出错，关键词: {keyword}")
            else:
                sys.stderr.write(f"ERROR: 搜索关键词 {keyword} 失败\n")
                sys.stderr.flush()
                self.log_handler.logger.error(f"搜索关键词 {keyword} 失败")
                if status_callback:
                    status_callback(f"搜索失败，关键词: {keyword}")
            
        except Exception as e:
            sys.stderr.write(f"ERROR: 采集过程发生错误: {str(e)}\n")
            sys.stderr.flush()
            self.log_handler.logger.error(f"采集过程发生错误: {str(e)}")
            if status_callback:
                status_callback(f"启动采集失败: {str(e)}")
        finally:
            # 完成后关闭浏览器
            if self.crawler:
                self.crawler.close()
                sys.stdout.write("INFO: 浏览器已关闭\n")
                sys.stdout.flush()
                self.log_handler.logger.info("浏览器已关闭")
            self.is_running = False
            # 恢复原始print函数
            self.restore_print()

    def start_crawler(self, notes_count, comments_count, keyword):
        """启动爬虫"""
        try:
            # 检查参数
            if not keyword.strip():
                return "请先添加关键词"
            
            if self.is_running:
                return "已有采集任务正在运行，请先停止"
            
            try:
                notes_count = int(notes_count)
                comments_count = int(comments_count)
            except:
                return "请输入有效的数字"
            
            if notes_count <= 0 or comments_count <= 0:
                return "采集数量必须大于0"
            
            # 检查cookie状态
            cookie_status = self.check_cookie_status()
            if "无效" in cookie_status or "错误" in cookie_status:
                return cookie_status
            
            # 使用线程运行爬虫，避免阻塞界面
            def status_updater(message):
                # 使用直接写入方式，避免递归
                sys.stdout.write(f"INFO: {message}\n")
                sys.stdout.flush()
                self.log_handler.logger.info(message)
            
            crawler_thread = threading.Thread(
                target=self.crawler_thread,
                args=(notes_count, comments_count, keyword, status_updater)
            )
            crawler_thread.daemon = True  # 设置为守护线程
            crawler_thread.start()
            
            return f"已启动采集任务，关键词: {keyword}"
            
        except ValueError:
            return "请输入有效的数字"
        except Exception as e:
            return f"启动采集失败: {str(e)}"

    def stop_crawler(self):
        """停止爬虫"""
        if not self.is_running:
            return "没有正在运行的采集任务"
        
        self.is_running = False
        if self.crawler:
            try:
                self.crawler.close()
                # 直接写入日志流
                sys.stdout.write("INFO: 已强制停止采集任务\n")
                sys.stdout.flush()
                self.log_handler.logger.info("已强制停止采集任务")
            except:
                pass
        
        return "采集任务已停止"
        
    def toggle_auto_refresh(self, state):
        """切换自动刷新状态"""
        self.auto_refresh_enabled = state
        
        # 如果启用自动刷新并且刷新线程未运行，则启动刷新线程
        if state and (self.refresh_thread is None or not self.refresh_thread.is_alive()):
            self.stop_refresh = False
            self.start_log_refresh_thread()
            
        # 如果禁用自动刷新且刷新线程正在运行，则停止刷新线程
        elif not state and self.refresh_thread and self.refresh_thread.is_alive():
            self.stop_refresh = True
            
        msg = "已启用日志自动刷新" if state else "已关闭日志自动刷新"
        return msg
        
    def start_log_refresh_thread(self):
        """启动日志刷新线程"""
        def refresh_worker():
            """日志刷新工作线程"""
            while not self.stop_refresh:
                time.sleep(1)  # 每秒更新一次
                if self.auto_refresh_enabled and self.logs_update_callback:
                    try:
                        # 尝试更新日志
                        self.logs_update_callback()
                    except:
                        pass  # 忽略异常，避免刷新线程崩溃
            
        self.refresh_thread = threading.Thread(
            target=refresh_worker,
            daemon=True
        )
        self.refresh_thread.start()

    def create_ui(self):
        """创建Gradio界面"""
        def refresh_logs_manually():
            """手动刷新日志的回调函数"""
            return self.get_logs()
        
        def run_crawler(notes_count, comments_count, keyword):
            """运行爬虫并返回状态和初始日志"""
            status = self.start_crawler(notes_count, comments_count, keyword)
            # 保存刷新回调函数引用
            self.logs_update_callback = lambda: auto_refresh_logs_output.update(value=self.get_logs())
            # 启动自动刷新线程
            if self.auto_refresh_enabled and (self.refresh_thread is None or not self.refresh_thread.is_alive()):
                self.start_log_refresh_thread()
            return status, self.get_logs()
        
        def stop_crawler_and_refresh():
            """停止爬虫并返回状态和日志"""
            status = self.stop_crawler()
            return status, self.get_logs()
        
        def clear_logs():
            """清空日志"""
            self.log_handler.clear()
            return ""
            
        # 创建更新函数引用
        auto_refresh_logs_output = None  # 占位符，将在UI创建后赋值
        
        with gr.Blocks(title="小红书数据采集工具", theme=gr.themes.Soft()) as interface:
            gr.Markdown("# 小红书数据采集工具 (实时日志版)")
            
            with gr.Tab("采集配置"):
                with gr.Row():
                    notes_count = gr.Number(
                        value=5,  # 默认改为较小的值，避免过多采集
                        label="采集笔记数量", 
                        minimum=1, 
                        maximum=100
                    )
                    comments_count = gr.Number(
                        value=20, 
                        label="每篇笔记采集评论数量", 
                        minimum=1, 
                        maximum=100
                    )
                
                with gr.Row():
                    keyword_input = gr.Textbox(
                        label="输入关键词", 
                        placeholder="请输入要采集的关键词"
                    )
                    add_btn = gr.Button("设置关键词")
                    clear_btn = gr.Button("清空关键词")
                
                with gr.Row():
                    cookie_check_btn = gr.Button("检查Cookie状态")
                    cookie_status = gr.Textbox(label="Cookie状态", interactive=False)
                
                keyword_display = gr.Textbox(
                    label="当前关键词", 
                    interactive=False
                )
                
                status_output = gr.Textbox(
                    label="状态信息", 
                    interactive=False
                )
                
                with gr.Row():
                    start_btn = gr.Button("开始采集", variant="primary")
                    stop_btn = gr.Button("停止采集", variant="secondary")
                    clear_logs_btn = gr.Button("清空日志")
                    refresh_logs_btn = gr.Button("手动刷新", size="sm")
                    auto_refresh_checkbox = gr.Checkbox(
                        label="启用自动刷新", 
                        value=True,
                        interactive=True
                    )
                
                # 添加日志输出面板
                logs_output = gr.Textbox(
                    label="采集日志", 
                    interactive=False,
                    lines=20,  # 增加显示行数
                    max_lines=20,  # 最大显示行数
                    show_copy_button=True,  # 添加复制按钮
                    autoscroll=True  # 启用自动滚动
                )
                
                # 保存logs_output引用，供自动刷新线程使用
                auto_refresh_logs_output = logs_output
                
                # 添加CSS自定义样式
                gr.HTML("""
                <style>
                .gradio-container .output-textbox {
                    height: 300px !important;
                    overflow-y: auto !important;
                    font-family: monospace !important;
                    white-space: pre-wrap !important;
                }
                </style>
                """)
                
                # 添加自动刷新的JavaScript代码
                if gr.__version__.startswith("5"):
                    # Gradio 5.x 版本使用HTML添加前端刷新功能
                    gr.HTML("""
                    <script>
                    let autoRefreshInterval;
                    
                    // 自动刷新函数
                    function setupAutoRefresh() {
                      // 清除现有的定时器
                      if (autoRefreshInterval) {
                        clearInterval(autoRefreshInterval);
                        autoRefreshInterval = null;
                      }
                      
                      // 检查是否启用了自动刷新
                      const checkbox = document.querySelector('input[type="checkbox"]');
                      if (checkbox && checkbox.checked) {
                        // 创建一个新的定时器
                        autoRefreshInterval = setInterval(() => {
                          // 查找并点击刷新按钮
                          const refreshButton = Array.from(document.querySelectorAll('button')).find(
                            btn => btn.textContent.includes('手动刷新')
                          );
                          if (refreshButton) {
                            refreshButton.click();
                          }
                        }, 2000); // 每2秒刷新一次
                      }
                    }
                    
                    // 页面加载后设置定时器
                    document.addEventListener('DOMContentLoaded', () => {
                      // 初始设置
                      setTimeout(setupAutoRefresh, 1000);
                      
                      // 监听复选框变化
                      const checkbox = document.querySelector('input[type="checkbox"]');
                      if (checkbox) {
                        checkbox.addEventListener('change', setupAutoRefresh);
                      }
                      
                      // 页面关闭前清理
                      window.addEventListener('beforeunload', () => {
                        if (autoRefreshInterval) {
                          clearInterval(autoRefreshInterval);
                        }
                      });
                    });
                    </script>
                    """)
                
                # 绑定事件 - 使用最基本的方式
                add_btn.click(
                    self.add_keyword,
                    inputs=[keyword_input],
                    outputs=[keyword_display, status_output]
                )
                
                clear_btn.click(
                    self.clear_keyword,
                    outputs=[keyword_display, status_output]
                )
                
                cookie_check_btn.click(
                    self.check_cookie_status,
                    outputs=cookie_status
                )
                
                # 开始采集 - 同时更新状态和日志
                start_btn.click(
                    run_crawler,
                    inputs=[notes_count, comments_count, keyword_display],
                    outputs=[status_output, logs_output]
                )
                
                # 停止采集 - 同时更新状态和日志
                stop_btn.click(
                    stop_crawler_and_refresh,
                    outputs=[status_output, logs_output]
                )
                
                # 清空日志
                clear_logs_btn.click(
                    clear_logs,
                    outputs=logs_output
                )
                
                # 手动刷新日志
                refresh_logs_btn.click(
                    refresh_logs_manually,
                    outputs=logs_output
                )
                
                # 切换自动刷新
                auto_refresh_checkbox.change(
                    self.toggle_auto_refresh,
                    inputs=[auto_refresh_checkbox],
                    outputs=[status_output]
                )
                
            # 添加使用说明标签
            with gr.Tab("使用说明"):
                gr.Markdown("""
                ## 使用步骤
                1. 在"采集配置"标签页设置采集参数：
                   - 设置要采集的笔记数量
                   - 设置每篇笔记要采集的评论数量
                2. 输入并设置要采集的关键词
                3. 点击"开始采集"启动采集任务
                4. 日志会自动刷新显示采集进度（也可手动刷新）
                
                ## 注意事项
                - 采集数量建议不要设置过大
                - 采集过程中可以随时停止
                - 如遇到采集失败，请检查Cookie是否有效
                - 如果日志刷新影响性能，可以关闭自动刷新功能
                """)
        
        # 保存UI引用，供刷新线程使用
        self.logs_update_callback = lambda: auto_refresh_logs_output.update(value=self.get_logs())
        
        # 启动日志刷新线程（如果启用了自动刷新）
        if self.auto_refresh_enabled:
            self.start_log_refresh_thread()
            
        return interface

def main(port=7860):
    print(f"启动小红书采集工具（实时日志版）在端口 {port}...")
    crawler_web = XHSCrawlerWeb()
    interface = crawler_web.create_ui()
    
    try:
        # 启动Gradio应用并自动打开浏览器
        # 先创建一个线程用于启动服务器
        server_thread = threading.Thread(
            target=interface.launch,
            kwargs={
                "server_name": "127.0.0.1",   # 本地访问
                "server_port": port,          # 设置端口
                "share": False,              # 不生成公共链接
                "prevent_thread_lock": True,  # 防止线程锁定
                "inbrowser": True             # 自动在浏览器中打开
            }
        )
        server_thread.daemon = True
        server_thread.start()
        
        # 主线程等待，让GUI保持运行
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print("程序已退出")
    except Exception as e:
        print(f"启动服务失败: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    # 检查是否有端口参数
    import sys
    port = 7860  # 默认端口
    if len(sys.argv) > 1:
        try:
            port = int(sys.argv[1])
        except ValueError:
            print(f"无效的端口号: {sys.argv[1]}，使用默认端口7860")
    
    main(port) 